import asyncio
import websockets
import pyaudio
import opuslib
import json
import signal
import sys
import wave
import time
import httpx
import requests
from pynput import keyboard as pynput_keyboard
import threading

MAC_ADDR = 'cd:62:f4:3d:b4:ff'
OTA_VERSION_URL = 'https://api.tenclass.net/xiaozhi/ota/'
# WebSocket 服务器地址
#SERVER_URI = 'wss://api.tenclass.net/xiaozhi/v1/'
#SERVER_URI = "ws://43.167.245.13:6699"
SERVER_URI = "ws://localhost:8082"
session_id = None
aes_opus_info = {"type": "hello", "version": 3, "transport": "websocket",
                 "audio_params": {"format": "opus", "sample_rate": 24000, "channels": 1, "frame_duration": 60},
                 "session_id": "b23ebfe9"}
# 音频参数
SAMPLE_RATE = 24000
CHANNELS = 1
FRAME_DURATION_MS = 60
SAMPLES_PER_FRAME = SAMPLE_RATE * FRAME_DURATION_MS // 1000
BUFFER_SIZE = SAMPLES_PER_FRAME * 2  # 16-bit samples

# 全局状态管理
tts_state = None
conn_state = False
wss = None
class AppState:
    def __init__(self):
        self.state = 'idle'
        self.lock = asyncio.Lock()

state = AppState()

# 创建 Opus 编解码器
encoder = opuslib.Encoder(SAMPLE_RATE, CHANNELS, opuslib.APPLICATION_AUDIO)
decoder = opuslib.Decoder(SAMPLE_RATE, CHANNELS)

# PyAudio 初始化
p = pyaudio.PyAudio()
input_stream = p.open(
    format=pyaudio.paInt16,
    channels=CHANNELS,
    rate=SAMPLE_RATE,
    input=True,
    frames_per_buffer=BUFFER_SIZE,
    start=False
)

output_stream = p.open(
    format=pyaudio.paInt16,
    channels=CHANNELS,
    rate=SAMPLE_RATE,
    output=True,
    frames_per_buffer=BUFFER_SIZE,
    start=False
)

async def send_hello_msg(ws):
    # 发送hello消息,建立websocket连接
    hello_message = {
        "type": "hello",
        "devicetype": "lichuang-dev",
        "version": 1,
        "transport": "websocket",
        "status": "logining",
        "macaddr": MAC_ADDR,
        "audio_params": {
            "format": "opus",
            "sample_rate": SAMPLE_RATE,
            "channels": CHANNELS,
            "frame_duration": FRAME_DURATION_MS
        }
    }
    await ws.send(json.dumps(hello_message))
    print("📨 初始化握手完成")
    print(f"send hello message: {hello_message}")

async def get_ota_version():
    header = {
        'Device-Id': MAC_ADDR,
        'Content-Type': 'application/json'
    }
    post_data = {"flash_size": 16777216, "minimum_free_heap_size": 8318916, "mac_address": f"{MAC_ADDR}",
                 "chip_model_name": "esp32s3", "chip_info": {"model": 9, "cores": 2, "revision": 2, "features": 18},
                 "application": {"name": "xiaozhi", "version": "0.9.9", "compile_time": "Jan 22 2025T20:40:23Z",
                                 "idf_version": "v5.3.2-dirty",
                                 "elf_sha256": "22986216df095587c42f8aeb06b239781c68ad8df80321e260556da7fcf5f522"},
                 "partition_table": [{"label": "nvs", "type": 1, "subtype": 2, "address": 36864, "size": 16384},
                                     {"label": "otadata", "type": 1, "subtype": 0, "address": 53248, "size": 8192},
                                     {"label": "phy_init", "type": 1, "subtype": 1, "address": 61440, "size": 4096},
                                     {"label": "model", "type": 1, "subtype": 130, "address": 65536, "size": 983040},
                                     {"label": "storage", "type": 1, "subtype": 130, "address": 1048576,
                                      "size": 1048576},
                                     {"label": "factory", "type": 0, "subtype": 0, "address": 2097152, "size": 4194304},
                                     {"label": "ota_0", "type": 0, "subtype": 16, "address": 6291456, "size": 4194304},
                                     {"label": "ota_1", "type": 0, "subtype": 17, "address": 10485760,
                                      "size": 4194304}],
                 "ota": {"label": "factory"},
                 "board": {"type": "lichuang-dev", "ssid": "mzy", "rssi": -58, "channel": 6,
                           "ip": "192.168.124.38", "mac": MAC_ADDR}}

    async with httpx.AsyncClient() as client:
        response = await client.post(OTA_VERSION_URL, headers=header, json=post_data)
        print('=========================')
        print(f'res:{response.text}')
        print(f"get version: {response}")

    # response = await requests.post(OTA_VERSION_URL, headers=header, data=json.dumps(post_data))
    # print('=========================')
    # print(f'res:{response.text}')
    # print(f"get version: {response}")


async def send_audio():
    global wss
    ws = wss
    """发送音频数据任务"""
    input_stream.start_stream()
    try:
        while True:
            async with state.lock:
                if state.state != "lisenting":
                    await asyncio.sleep(0.03)
                    continue

            try:
                pcm_data = input_stream.read(BUFFER_SIZE, exception_on_overflow=False)
                opus_packet = encoder.encode(pcm_data, BUFFER_SIZE)
                await ws.send(opus_packet)
                #print(f"📤 发送 Opus 数据，大小: {len(opus_packet)} 字节")
            except (websockets.exceptions.ConnectionClosed, OSError) as e:
                print(f"❌ 发送音频错误: {e}")
                break
            except Exception as e:
                print(f"⚠️ 编码错误: {e}")
    finally:
        input_stream.stop_stream()

async def receive_messages():
    """接收消息任务"""
    global tts_state,wss
    ws = wss
    output_stream.start_stream()
    try:
        async for message in ws:
            try:
                if isinstance(message, str):
                    message_data = json.loads(message)
                    print(f'📥 接收到控制消息: {message_data}')
                    async with state.lock:
                        if message_data.get("type") == "hello":
                            state.state = 'ilde'
                            session_id = message_data.get("session_id")
                            print("\033[32m Set state ilde\033[0m")  # 绿色
                            print(f'session_id: {session_id}')
                        elif message_data.get("type") == "tts":
                            if message_data.get("state") == 'start':
                                state.state = 'speaking'
                                tts_state = "start"
                                print("\033[32m Set state speaking\033[0m")  # 绿色
                            elif message_data.get("state") == 'sentence_start':
                                print(f'Speak Text:  {message_data.get("text")}')
                            elif message_data.get("state") == 'stop':
                                state.state = 'idle'
                                lisening_message = {"session_id": session_id,"type":"listen","state":"start","mode":"auto"}
                                await ws.send(json.dumps(lisening_message))
                                print("\033[32m Set state lisenting\033[0m")  # 绿色
                            
                elif isinstance(message, bytes):
                    #async with state.lock:
                    if state.state == 'speaking':
                        decoded_pcm = decoder.decode(message, BUFFER_SIZE)
                        output_stream.write(decoded_pcm)
                        #print(f"🔊 播放音频，大小: {int(time.time() * 1000)} {len(message)} 字节")
            except Exception as e:
                print(f"⚠️ 处理消息错误: {e}")
    finally:
        output_stream.stop_stream()

async def user_input_handler(ws):
    """处理用户输入"""
    loop = asyncio.get_running_loop()
    while True:
        try:
            user_input = await loop.run_in_executor(None, input, "输入消息 (输入 exit 退出): ")
            if user_input == '你好小智' and state.state == 'idle':
                detect_message = {"session_id":session_id, "type":"listen","state":"detect","text":"你好小智"}
                await ws.send(json.dumps(detect_message))
                continue
            if user_input.lower() == "exit":
                await ws.close()
                break

            message = {"type": "chat", "content": user_input}
            #message = {"type": "tts", "content": user_input}
            await ws.send(json.dumps(message))
        except (asyncio.CancelledError, KeyboardInterrupt):
            break

def send_audio_thread_function(loop):
    asyncio.run_coroutine_threadsafe(send_audio(), loop)
def recv_message_thread_function(loop):
    asyncio.run_coroutine_threadsafe(receive_messages(), loop)

async def websocket_client():
    """WebSocket 客户端主函数"""
    global conn_state, wss
    async with websockets.connect(SERVER_URI) as ws:
        conn_state = True
        wss = ws
        # 启动一个线程，用于接收音频数据
        loop = asyncio.get_running_loop()
        thread1 = threading.Thread(target=recv_message_thread_function, args=(loop,))
        thread1.start()
        thread2 = threading.Thread(target=send_audio_thread_function, args=(loop,))
        thread2.start()
        # 启动所有任务
        tasks = [
            asyncio.create_task(user_input_handler(ws))
        ]
        try:
            await send_hello_msg(ws)
            await asyncio.gather(*tasks)
        except websockets.exceptions.ConnectionClosed:
            print("❌ 连接已关闭")
            conn_state = False
        finally:
            print('取消所有任务')
            # 取消所有任务
            for task in tasks:
                task.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)

def cleanup(signal, frame):
    """清理资源"""
    print("\n🛑 正在清理资源...")
    input_stream.close()
    output_stream.close()
    p.terminate()
    sys.exit(0)

def on_space_key_press(event):
    global key_state, wss, aes_opus_info, conn_state
    if key_state == "press":
        return
    key_state = "press"
    # 判断是否需要发送hello消息
    if conn_state is False or aes_opus_info['session_id'] is None:
        conn_state = True
        send_hello_msg(wss)
    if tts_state == "start" or tts_state == "entence_start":
        # 在播放状态下发送abort消息
        wss.send(json.dumps({"type": "abort"}))
        print(f"send abort message")
    if aes_opus_info['session_id'] is not None:
        # 发送start listen消息
        msg = {"session_id": aes_opus_info['session_id'], "type": "listen", "state": "start", "mode": "manual"}
        print(f"send start listen message: {msg}")
        wss.send(json.dumps(msg))


def on_space_key_release(event):
    global aes_opus_info, key_state
    key_state = "release"
    # 发送stop listen消息
    if aes_opus_info['session_id'] is not None:
        msg = {"session_id": aes_opus_info['session_id'], "type": "listen", "state": "stop"}
        print(f"send stop listen message: {msg}")
        wss.send(json.dumps(msg))

def on_press(key):
    if key == pynput_keyboard.Key.space:
        on_space_key_press(None)


def on_release(key):
    if key == pynput_keyboard.Key.space:
        on_space_key_release(None)
    # Stop listener
    if key == pynput_keyboard.Key.esc:
        return False

async def main():
    # 注册信号处理
    signal.signal(signal.SIGINT, cleanup)
    signal.signal(signal.SIGTERM, cleanup)

    # 调用 get_ota_version()，使用 await
    await get_ota_version()

    try:
        # 监听键盘按键，当按下空格键时，发送 listen 消息
        # listener = pynput_keyboard.Listener(on_press=on_press, on_release=on_release)
        # listener.start()
        await websocket_client()
    finally:
        print("finished")
        cleanup(None, None)

if __name__ == "__main__":
    asyncio.run(main())  